import numpy as np
import skimage.io as imgio
import skimage
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_labels

h,w = 321,321
categorys_num = 21
crf_config = {"bi_sxy":50,"bi_srgb":10,"bi_compat":5,"g_sxy":3,"g_compat":3,"iterations":5}
img_mean = np.ones((h,w,3))
img_mean[:,:,0] *= 104.00698793
img_mean[:,:,1] *= 116.66876762
img_mean[:,:,2] *= 122.67891434
def crf_inference(label,img):
    '''
    label: the predicted label map of cnn, shape [h,w]
    img: the origin img, shape [h,w,3]
    '''
    crf = dcrf.DenseCRF2D(w,h,categorys_num)

    unary = unary_from_labels(label, categorys_num, 0.9, zero_unsure=False)

    crf.setUnaryEnergy(unary)

    # pairwise energy
    crf.addPairwiseGaussian( sxy=crf_config["g_sxy"], compat=crf_config["g_compat"] )
    #print("img shape:%s" % str(img.shape))
    crf.addPairwiseBilateral( sxy=crf_config["bi_sxy"], srgb=crf_config["bi_srgb"], rgbim=img, compat=crf_config["bi_compat"] )
    Q = crf.inference( crf_config["iterations"] )
    #np.save("Q.npy",Q)
    r = np.argmax(Q,axis=0).reshape((h,w))
    return r

if __name__ == "__main__":
    rate = int(250 / 21)
    file_path = "test"
    img = imgio.imread("%s.png" % file_path)
    img = skimage.img_as_ubyte(img)
    predicted_label = imgio.imread("%s_label.png" % file_path) / rate
    predicted_label = predicted_label.astype(np.uint8)
    crf_img = crf_inference(predicted_label,img)
    crf_img = crf_img.astype(np.uint8)*rate
    imgio.imsave("pred.png" ,crf_img)
